{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 隐语SecretFlow实际场景MPC算法开发实践\n",
    "\n",
    "> This tutorial is only available in Chinese.\n",
    "\n",
    "推荐使用`conda`创建一个新环境\n",
    "> conda create -n sf python=3.8\n",
    "\n",
    "直接使用`pip`安装secretflow\n",
    "> pip install -U secretflow\n",
    "\n",
    "基于secretflow：**0.8.2b2**版本\n",
    "\n",
    "此代码示例主要是展示了如何基于secretflow以及SPU隐私计算设备完成一个实际的应用的开发，推荐先看前一个教程[spu_basics](./spu_basics.ipynb)熟悉基本的SPU概念。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 任务介绍"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Vehicle Insurance Claim Fraud Detection\n",
    "\n",
    "该数据集来源于[kaggle](https://www.kaggle.com/datasets/shivamb/vehicle-claim-fraud-detection)，包含\n",
    "- 车辆数据集-属性、模型、事故详细信息等\n",
    "- 保单详细信息-保单类型、有效期等\n",
    "\n",
    "目标是检测索赔申请是否欺诈： 字段`FraudFound_P` (0 or 1) 即为预测的target值，是一个典型的**二分类场景**。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 实验目标\n",
    "在本次实验中，我们将会利用一个开源数据集在隐语上完成隐私保护的逻辑回归、神经网络模型和XGB模型。主要涉及到如下的几个流程：\n",
    "1. 数据加载\n",
    "2. 数据洞察\n",
    "3. 数据预处理\n",
    "4. 模型构建\n",
    "5. 模型的训练与预测"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 前置工作"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ray集群启动（多机部署）\n",
    "考虑多机部署的情况，在启动secretflow之前需要先将ray集群启动。在header节点和worker节点上各自执行下述的指令。\n",
    "> P.S. 启动集群之后，可以执行`ray status`看一下集群是否正确启动完成"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Header节点**\n",
    "``` Bash\n",
    "RAY_DISABLE_REMOTE_CODE=true \\\n",
    "ray start --head --node-ip-address=\"head_ip\" --port=\"head_port\" --resources='{\"alice\": 20}' --include-dashboard=False\n",
    "```\n",
    "**Worker节点**\n",
    "``` Bash\n",
    "RAY_DISABLE_REMOTE_CODE=true \\\n",
    "ray start --address=\"head_ip:head_port\" --resources='{\"bob\": 20}'\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 如下是多机版初始化secretflow的代码，需要给出header节点的IP和PORT\n",
    "# head_ip = \"xxx\"\n",
    "# head_port = \"xxx\" \n",
    "# sf.init(address=f'{head_ip}:{head_port}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 单机部署\n",
    "我们在此使用单机部署的方式做一个样例展示。\n",
    "通过调用`sf.init()`我们实例化了一个ray集群，有5个节点，也就对应了5个物理设备。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-05-24 14:55:43,444\tINFO worker.py:1544 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265 \u001b[39m\u001b[22m\n"
     ]
    }
   ],
   "source": [
    "import secretflow as sf\n",
    "# Check the version of your SecretFlow\n",
    "print('The version of SecretFlow: {}'.format(sf.__version__))\n",
    "\n",
    "sf.shutdown()\n",
    "# Standalone Mode\n",
    "sf.init(\n",
    "    ['alice', 'bob', 'carol', 'davy', 'eric'],\n",
    "    address='local'\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义明文计算设备PYU\n",
    "我们在启动了上述5个节点之后，明确隐语中的逻辑设备。这里我们将alice、bob、carol三方作为数据的提供方，可以本地执行明文计算，也就是 **PYU (PYthon runtime Unit)** 设备。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "alice\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[33m(raylet)\u001b[0m /home/zoupeicheng.zpc/miniconda3/envs/py38/lib/python3.8/site-packages/ray/dashboard/modules/reporter/reporter_agent.py:56: UserWarning: `gpustat` package is not installed. GPU monitoring is not available. To have full functionality of the dashboard please install `pip install ray[default]`.)\n",
      "\u001b[2m\u001b[33m(raylet)\u001b[0m   warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "alice = sf.PYU('alice')\n",
    "bob = sf.PYU('bob')\n",
    "carol = sf.PYU('carol')\n",
    "\n",
    "print(alice)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 定义密文计算设备SPU (3PC)\n",
    "进一步，我们以**SPU (Secure Processing Unit)** 为例，选择3个物理节点组成基于MPC（下例为三方的ABY3协议）的隐私计算设备。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import spu\n",
    "from secretflow.utils.testing import unused_tcp_port\n",
    "\n",
    "aby3_cluster_def = {\n",
    "    'nodes': [\n",
    "        {\n",
    "            'party': 'alice',\n",
    "            'address': f'127.0.0.1:{unused_tcp_port()}',\n",
    "        },\n",
    "        {'party': 'bob', 'id': 'local:1', 'address': f'127.0.0.1:{unused_tcp_port()}'},\n",
    "        {\n",
    "            'party': 'carol',\n",
    "            'address': f'127.0.0.1:{unused_tcp_port()}',\n",
    "        },\n",
    "    ],\n",
    "    'runtime_config': {\n",
    "        'protocol': spu.spu_pb2.ABY3,\n",
    "        'field': spu.spu_pb2.FM64,\n",
    "    },\n",
    "}\n",
    "\n",
    "my_spu = sf.SPU(aby3_cluster_def)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据加载"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data (Mock)\n",
    "在定义好隐语中的逻辑设备概念之后，我们演示一下如何进行数据的读入。这里使用一个mock的data load方法`get_data_mock()`做一个演示。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x_plaintext: 2\n"
     ]
    }
   ],
   "source": [
    "def get_data_mock():\n",
    "    return 2\n",
    "\n",
    "x_plaintext = get_data_mock()\n",
    "print(f\"x_plaintext: {x_plaintext}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "指定PYU设备读取数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Plaintext Python Object: 2, PYU object: <secretflow.device.device.pyu.PYUObject object at 0x7fcfba18db20>\n",
      "Reveal PYU object: 2\n"
     ]
    }
   ],
   "source": [
    "x_alice_pyu = alice(get_data_mock)()\n",
    "\n",
    "print(f\"Plaintext Python Object: {x_plaintext}, PYU object: {x_alice_pyu}\")\n",
    "print(f\"Reveal PYU object: {sf.reveal(x_alice_pyu)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PYU->SPU 数据转换"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SPU object: <secretflow.device.device.spu.SPUObject object at 0x7fd000234ca0>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=2294419)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: \n",
      "\u001b[2m\u001b[36m(_run pid=2294419)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294419)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294419)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=2294419)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=2294419)\u001b[0m WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Reveal SPU object: 2\n"
     ]
    }
   ],
   "source": [
    "x_alice_spu = x_alice_pyu.to(my_spu)\n",
    "print(f\"SPU object: {x_alice_spu}\")\n",
    "\n",
    "print(f\"Reveal SPU object: {sf.reveal(x_alice_spu)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data (Distributed)\n",
    "我们下面考虑对一个实际应用场景的数据进行读取，也就是全集数据垂直分布在不同的参与方中。\n",
    "> 出于演示的目的，我们这里将中心化的明文数据进行垂直分割的拆分，首先观察下此数据集的特征。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 读入明文全集数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "File already downloaded.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "\"\"\"\n",
    "Create dir to save dataset files\n",
    "This will create a directory `data` to store the dataset file\n",
    "\"\"\"\n",
    "if not os.path.exists('data'):\n",
    "    os.mkdir('data')\n",
    "\n",
    "\"\"\"\n",
    "The original data is from Kaggle: https://www.kaggle.com/datasets/shivamb/vehicle-claim-fraud-detection. \n",
    "We promise we only use the data for demo only.\n",
    "\"\"\"\n",
    "path = \"https://secretflow-data.oss-accelerate.aliyuncs.com/datasets/vehicle_nsurance_claim/fraud_oracle.csv\"\n",
    "if not os.path.exists('data/fraud_oracle.csv'):\n",
    "    res = os.system('cd data && wget {}'.format(path))\n",
    "    if res != 0:\n",
    "        raise Exception('File: {} download fails!'.format(path))\n",
    "else:\n",
    "    print(f'File already downloaded.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Month</th>\n",
       "      <th>WeekOfMonth</th>\n",
       "      <th>DayOfWeek</th>\n",
       "      <th>Make</th>\n",
       "      <th>AccidentArea</th>\n",
       "      <th>DayOfWeekClaimed</th>\n",
       "      <th>MonthClaimed</th>\n",
       "      <th>WeekOfMonthClaimed</th>\n",
       "      <th>Sex</th>\n",
       "      <th>MaritalStatus</th>\n",
       "      <th>...</th>\n",
       "      <th>AgeOfVehicle</th>\n",
       "      <th>AgeOfPolicyHolder</th>\n",
       "      <th>PoliceReportFiled</th>\n",
       "      <th>WitnessPresent</th>\n",
       "      <th>AgentType</th>\n",
       "      <th>NumberOfSuppliments</th>\n",
       "      <th>AddressChange_Claim</th>\n",
       "      <th>NumberOfCars</th>\n",
       "      <th>Year</th>\n",
       "      <th>BasePolicy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Dec</td>\n",
       "      <td>5</td>\n",
       "      <td>Wednesday</td>\n",
       "      <td>Honda</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Tuesday</td>\n",
       "      <td>Jan</td>\n",
       "      <td>1</td>\n",
       "      <td>Female</td>\n",
       "      <td>Single</td>\n",
       "      <td>...</td>\n",
       "      <td>3 years</td>\n",
       "      <td>26 to 30</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>none</td>\n",
       "      <td>1 year</td>\n",
       "      <td>3 to 4</td>\n",
       "      <td>1994</td>\n",
       "      <td>Liability</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Jan</td>\n",
       "      <td>3</td>\n",
       "      <td>Wednesday</td>\n",
       "      <td>Honda</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Monday</td>\n",
       "      <td>Jan</td>\n",
       "      <td>4</td>\n",
       "      <td>Male</td>\n",
       "      <td>Single</td>\n",
       "      <td>...</td>\n",
       "      <td>6 years</td>\n",
       "      <td>31 to 35</td>\n",
       "      <td>Yes</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>none</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1994</td>\n",
       "      <td>Collision</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Oct</td>\n",
       "      <td>5</td>\n",
       "      <td>Friday</td>\n",
       "      <td>Honda</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Thursday</td>\n",
       "      <td>Nov</td>\n",
       "      <td>2</td>\n",
       "      <td>Male</td>\n",
       "      <td>Married</td>\n",
       "      <td>...</td>\n",
       "      <td>7 years</td>\n",
       "      <td>41 to 50</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>none</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1994</td>\n",
       "      <td>Collision</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Jun</td>\n",
       "      <td>2</td>\n",
       "      <td>Saturday</td>\n",
       "      <td>Toyota</td>\n",
       "      <td>Rural</td>\n",
       "      <td>Friday</td>\n",
       "      <td>Jul</td>\n",
       "      <td>1</td>\n",
       "      <td>Male</td>\n",
       "      <td>Married</td>\n",
       "      <td>...</td>\n",
       "      <td>more than 7</td>\n",
       "      <td>51 to 65</td>\n",
       "      <td>Yes</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>more than 5</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1994</td>\n",
       "      <td>Liability</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Jan</td>\n",
       "      <td>5</td>\n",
       "      <td>Monday</td>\n",
       "      <td>Honda</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Tuesday</td>\n",
       "      <td>Feb</td>\n",
       "      <td>2</td>\n",
       "      <td>Female</td>\n",
       "      <td>Single</td>\n",
       "      <td>...</td>\n",
       "      <td>5 years</td>\n",
       "      <td>31 to 35</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>none</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1994</td>\n",
       "      <td>Collision</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 33 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "  Month  WeekOfMonth  DayOfWeek    Make AccidentArea DayOfWeekClaimed  \\\n",
       "0   Dec            5  Wednesday   Honda        Urban          Tuesday   \n",
       "1   Jan            3  Wednesday   Honda        Urban           Monday   \n",
       "2   Oct            5     Friday   Honda        Urban         Thursday   \n",
       "3   Jun            2   Saturday  Toyota        Rural           Friday   \n",
       "4   Jan            5     Monday   Honda        Urban          Tuesday   \n",
       "\n",
       "  MonthClaimed  WeekOfMonthClaimed     Sex MaritalStatus  ...  AgeOfVehicle  \\\n",
       "0          Jan                   1  Female        Single  ...       3 years   \n",
       "1          Jan                   4    Male        Single  ...       6 years   \n",
       "2          Nov                   2    Male       Married  ...       7 years   \n",
       "3          Jul                   1    Male       Married  ...   more than 7   \n",
       "4          Feb                   2  Female        Single  ...       5 years   \n",
       "\n",
       "  AgeOfPolicyHolder PoliceReportFiled WitnessPresent AgentType  \\\n",
       "0          26 to 30                No             No  External   \n",
       "1          31 to 35               Yes             No  External   \n",
       "2          41 to 50                No             No  External   \n",
       "3          51 to 65               Yes             No  External   \n",
       "4          31 to 35                No             No  External   \n",
       "\n",
       "   NumberOfSuppliments  AddressChange_Claim  NumberOfCars  Year  BasePolicy  \n",
       "0                 none               1 year        3 to 4  1994   Liability  \n",
       "1                 none            no change     1 vehicle  1994   Collision  \n",
       "2                 none            no change     1 vehicle  1994   Collision  \n",
       "3          more than 5            no change     1 vehicle  1994   Liability  \n",
       "4                 none            no change     1 vehicle  1994   Collision  \n",
       "\n",
       "[5 rows x 33 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "import pandas as pd\n",
    "\n",
    "\"\"\"\n",
    "This should point to the data downloaded from Kaggle.\n",
    "By default, the .csv file shall be in the data directory\n",
    "\"\"\"\n",
    "full_data_path = 'data/fraud_oracle.csv'\n",
    "df = pd.read_csv(full_data_path)\n",
    "df.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 数据三方垂直拆分\n",
    "我们首先对这个数据进行一个拆分的处理，来模拟一个数据垂直分割的三方场景：\n",
    "\n",
    "- alice持有前10个属性\n",
    "- bob持有中间的10个属性\n",
    "- carol持有剩下的所有属性以及标签值\n",
    "\n",
    "同时为了方便各方之间的样本做对齐，我们加了一个新的特征`UID`来标识数据样本。\n",
    "\n",
    "我们预先基于sklearn将全集数据拆分成训练集和测试集，方便后续进行模型训练效果的验证。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(10792, 32)\n",
      "(4626, 32)\n"
     ]
    }
   ],
   "source": [
    "train_alice_path = \"data/alice_train.csv\"\n",
    "train_bob_path = \"data/bob_train.csv\"\n",
    "train_carol_path = \"data/carol_train.csv\"\n",
    "\n",
    "test_alice_path = \"data/alice_test.csv\"\n",
    "test_bob_path = \"data/bob_test.csv\"\n",
    "test_carol_path = \"data/carol_test.csv\"\n",
    "\n",
    "def load_dataset_full(data_path):\n",
    "    df = pd.read_csv(data_path)\n",
    "    df = df.drop([0])\n",
    "    df = df.loc[df['DayOfWeekClaimed']!='0']\n",
    "    y = df['FraudFound_P']\n",
    "    X = df.drop(columns='FraudFound_P')\n",
    "    return X, y\n",
    "\n",
    "def split_data():\n",
    "    x, y = load_dataset_full(full_data_path)\n",
    "    x_train, x_test, y_train, y_test = train_test_split(\n",
    "        x, y, test_size=0.3, random_state=10\n",
    "    )\n",
    "\n",
    "    print(x_train.shape)\n",
    "    train_alice_csv = x_train.iloc[:, :10]\n",
    "    train_bob_csv = x_train.iloc[:, 10:20]\n",
    "    train_carol_csv = pd.concat([x_train.iloc[:, 20:], y_train], axis=1)\n",
    "\n",
    "    train_alice_csv.to_csv(train_alice_path, index_label='UID')\n",
    "    train_bob_csv.to_csv(train_bob_path, index_label='UID')\n",
    "    train_carol_csv.to_csv(train_carol_path, index_label='UID')\n",
    "\n",
    "    print(x_test.shape)\n",
    "    test_alice_csv = x_test.iloc[:, :10]\n",
    "    test_bob_csv = x_test.iloc[:, 10:20]\n",
    "    test_carol_csv = pd.concat([x_test.iloc[:, 20:], y_test], axis=1)\n",
    "\n",
    "    test_alice_csv.to_csv(test_alice_path, index_label='UID')\n",
    "    test_bob_csv.to_csv(test_bob_path, index_label='UID')\n",
    "    test_carol_csv.to_csv(test_carol_path, index_label='UID')\n",
    "\n",
    "split_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>UID</th>\n",
       "      <th>Month</th>\n",
       "      <th>WeekOfMonth</th>\n",
       "      <th>DayOfWeek</th>\n",
       "      <th>Make</th>\n",
       "      <th>AccidentArea</th>\n",
       "      <th>DayOfWeekClaimed</th>\n",
       "      <th>MonthClaimed</th>\n",
       "      <th>WeekOfMonthClaimed</th>\n",
       "      <th>Sex</th>\n",
       "      <th>MaritalStatus</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2853</td>\n",
       "      <td>Mar</td>\n",
       "      <td>4</td>\n",
       "      <td>Sunday</td>\n",
       "      <td>Toyota</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Friday</td>\n",
       "      <td>Apr</td>\n",
       "      <td>1</td>\n",
       "      <td>Male</td>\n",
       "      <td>Married</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7261</td>\n",
       "      <td>Apr</td>\n",
       "      <td>4</td>\n",
       "      <td>Saturday</td>\n",
       "      <td>Honda</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Monday</td>\n",
       "      <td>Apr</td>\n",
       "      <td>4</td>\n",
       "      <td>Male</td>\n",
       "      <td>Married</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>9862</td>\n",
       "      <td>Jun</td>\n",
       "      <td>4</td>\n",
       "      <td>Sunday</td>\n",
       "      <td>Toyota</td>\n",
       "      <td>Rural</td>\n",
       "      <td>Monday</td>\n",
       "      <td>Jun</td>\n",
       "      <td>4</td>\n",
       "      <td>Female</td>\n",
       "      <td>Single</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>14037</td>\n",
       "      <td>Mar</td>\n",
       "      <td>2</td>\n",
       "      <td>Monday</td>\n",
       "      <td>Mazda</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Monday</td>\n",
       "      <td>Mar</td>\n",
       "      <td>2</td>\n",
       "      <td>Male</td>\n",
       "      <td>Single</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>10199</td>\n",
       "      <td>Jun</td>\n",
       "      <td>3</td>\n",
       "      <td>Friday</td>\n",
       "      <td>Mazda</td>\n",
       "      <td>Urban</td>\n",
       "      <td>Tuesday</td>\n",
       "      <td>Jun</td>\n",
       "      <td>4</td>\n",
       "      <td>Female</td>\n",
       "      <td>Single</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     UID Month  WeekOfMonth DayOfWeek    Make AccidentArea DayOfWeekClaimed  \\\n",
       "0   2853   Mar            4    Sunday  Toyota        Urban           Friday   \n",
       "1   7261   Apr            4  Saturday   Honda        Urban           Monday   \n",
       "2   9862   Jun            4    Sunday  Toyota        Rural           Monday   \n",
       "3  14037   Mar            2    Monday   Mazda        Urban           Monday   \n",
       "4  10199   Jun            3    Friday   Mazda        Urban          Tuesday   \n",
       "\n",
       "  MonthClaimed  WeekOfMonthClaimed     Sex MaritalStatus  \n",
       "0          Apr                   1    Male       Married  \n",
       "1          Apr                   4    Male       Married  \n",
       "2          Jun                   4  Female        Single  \n",
       "3          Mar                   2    Male        Single  \n",
       "4          Jun                   4  Female        Single  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alice_train_df = pd.read_csv(train_alice_path)\n",
    "alice_train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>UID</th>\n",
       "      <th>Age</th>\n",
       "      <th>Fault</th>\n",
       "      <th>PolicyType</th>\n",
       "      <th>VehicleCategory</th>\n",
       "      <th>VehiclePrice</th>\n",
       "      <th>PolicyNumber</th>\n",
       "      <th>RepNumber</th>\n",
       "      <th>Deductible</th>\n",
       "      <th>DriverRating</th>\n",
       "      <th>Days_Policy_Accident</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2853</td>\n",
       "      <td>39</td>\n",
       "      <td>Policy Holder</td>\n",
       "      <td>Sedan - All Perils</td>\n",
       "      <td>Sedan</td>\n",
       "      <td>20000 to 29000</td>\n",
       "      <td>2854</td>\n",
       "      <td>8</td>\n",
       "      <td>400</td>\n",
       "      <td>2</td>\n",
       "      <td>more than 30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7261</td>\n",
       "      <td>58</td>\n",
       "      <td>Policy Holder</td>\n",
       "      <td>Sedan - Liability</td>\n",
       "      <td>Sport</td>\n",
       "      <td>20000 to 29000</td>\n",
       "      <td>7262</td>\n",
       "      <td>4</td>\n",
       "      <td>400</td>\n",
       "      <td>4</td>\n",
       "      <td>more than 30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>9862</td>\n",
       "      <td>28</td>\n",
       "      <td>Policy Holder</td>\n",
       "      <td>Sedan - All Perils</td>\n",
       "      <td>Sedan</td>\n",
       "      <td>less than 20000</td>\n",
       "      <td>9863</td>\n",
       "      <td>5</td>\n",
       "      <td>400</td>\n",
       "      <td>4</td>\n",
       "      <td>more than 30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>14037</td>\n",
       "      <td>28</td>\n",
       "      <td>Policy Holder</td>\n",
       "      <td>Sedan - Collision</td>\n",
       "      <td>Sedan</td>\n",
       "      <td>20000 to 29000</td>\n",
       "      <td>14038</td>\n",
       "      <td>11</td>\n",
       "      <td>400</td>\n",
       "      <td>4</td>\n",
       "      <td>more than 30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>10199</td>\n",
       "      <td>35</td>\n",
       "      <td>Policy Holder</td>\n",
       "      <td>Sedan - Collision</td>\n",
       "      <td>Sedan</td>\n",
       "      <td>20000 to 29000</td>\n",
       "      <td>10200</td>\n",
       "      <td>12</td>\n",
       "      <td>400</td>\n",
       "      <td>4</td>\n",
       "      <td>more than 30</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     UID  Age          Fault          PolicyType VehicleCategory  \\\n",
       "0   2853   39  Policy Holder  Sedan - All Perils           Sedan   \n",
       "1   7261   58  Policy Holder   Sedan - Liability           Sport   \n",
       "2   9862   28  Policy Holder  Sedan - All Perils           Sedan   \n",
       "3  14037   28  Policy Holder   Sedan - Collision           Sedan   \n",
       "4  10199   35  Policy Holder   Sedan - Collision           Sedan   \n",
       "\n",
       "      VehiclePrice  PolicyNumber  RepNumber  Deductible  DriverRating  \\\n",
       "0   20000 to 29000          2854          8         400             2   \n",
       "1   20000 to 29000          7262          4         400             4   \n",
       "2  less than 20000          9863          5         400             4   \n",
       "3   20000 to 29000         14038         11         400             4   \n",
       "4   20000 to 29000         10200         12         400             4   \n",
       "\n",
       "  Days_Policy_Accident  \n",
       "0         more than 30  \n",
       "1         more than 30  \n",
       "2         more than 30  \n",
       "3         more than 30  \n",
       "4         more than 30  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bob_train_df = pd.read_csv(train_bob_path)\n",
    "bob_train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>UID</th>\n",
       "      <th>Days_Policy_Claim</th>\n",
       "      <th>PastNumberOfClaims</th>\n",
       "      <th>AgeOfVehicle</th>\n",
       "      <th>AgeOfPolicyHolder</th>\n",
       "      <th>PoliceReportFiled</th>\n",
       "      <th>WitnessPresent</th>\n",
       "      <th>AgentType</th>\n",
       "      <th>NumberOfSuppliments</th>\n",
       "      <th>AddressChange_Claim</th>\n",
       "      <th>NumberOfCars</th>\n",
       "      <th>Year</th>\n",
       "      <th>BasePolicy</th>\n",
       "      <th>FraudFound_P</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2853</td>\n",
       "      <td>more than 30</td>\n",
       "      <td>1</td>\n",
       "      <td>7 years</td>\n",
       "      <td>36 to 40</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>more than 5</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1994</td>\n",
       "      <td>All Perils</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>7261</td>\n",
       "      <td>more than 30</td>\n",
       "      <td>none</td>\n",
       "      <td>more than 7</td>\n",
       "      <td>51 to 65</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>1 to 2</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1995</td>\n",
       "      <td>Liability</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>9862</td>\n",
       "      <td>more than 30</td>\n",
       "      <td>none</td>\n",
       "      <td>7 years</td>\n",
       "      <td>31 to 35</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>none</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1995</td>\n",
       "      <td>All Perils</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>14037</td>\n",
       "      <td>more than 30</td>\n",
       "      <td>1</td>\n",
       "      <td>6 years</td>\n",
       "      <td>31 to 35</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>External</td>\n",
       "      <td>none</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1996</td>\n",
       "      <td>Collision</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>10199</td>\n",
       "      <td>more than 30</td>\n",
       "      <td>none</td>\n",
       "      <td>5 years</td>\n",
       "      <td>31 to 35</td>\n",
       "      <td>No</td>\n",
       "      <td>No</td>\n",
       "      <td>Internal</td>\n",
       "      <td>none</td>\n",
       "      <td>no change</td>\n",
       "      <td>1 vehicle</td>\n",
       "      <td>1995</td>\n",
       "      <td>Collision</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     UID Days_Policy_Claim PastNumberOfClaims AgeOfVehicle AgeOfPolicyHolder  \\\n",
       "0   2853      more than 30                  1      7 years          36 to 40   \n",
       "1   7261      more than 30               none  more than 7          51 to 65   \n",
       "2   9862      more than 30               none      7 years          31 to 35   \n",
       "3  14037      more than 30                  1      6 years          31 to 35   \n",
       "4  10199      more than 30               none      5 years          31 to 35   \n",
       "\n",
       "  PoliceReportFiled WitnessPresent AgentType NumberOfSuppliments  \\\n",
       "0                No             No  External         more than 5   \n",
       "1                No             No  External              1 to 2   \n",
       "2                No             No  External                none   \n",
       "3                No             No  External                none   \n",
       "4                No             No  Internal                none   \n",
       "\n",
       "  AddressChange_Claim NumberOfCars  Year  BasePolicy  FraudFound_P  \n",
       "0           no change    1 vehicle  1994  All Perils             0  \n",
       "1           no change    1 vehicle  1995   Liability             0  \n",
       "2           no change    1 vehicle  1995  All Perils             0  \n",
       "3           no change    1 vehicle  1996   Collision             0  \n",
       "4           no change    1 vehicle  1995   Collision             0  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "carol_train_df = pd.read_csv(train_carol_path)\n",
    "carol_train_df.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 三方数据加载\n",
    "注意：这里的接口里面需要显示地指明用于多方之间样本对齐的key，以及明确使用何种设备来执行PSI。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.244 [info] [bucket_psi.cc:Init:228] bucket size set to 1048576\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.269 [info] [bucket_psi.cc:Run:97] Begin sanity check for input file: data/carol_train.csv, precheck_switch:true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.290 [info] [csv_checker.cc:CsvChecker:121] Executing duplicated scripts: LC_ALL=C sort --buffer-size=1G --temporary-directory=data --stable selected-keys.1684911350271260403 | LC_ALL=C uniq -d > duplicate-keys.1684911350271260403\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.320 [info] [bucket_psi.cc:Run:115] End sanity check for input file: data/carol_train.csv, size=10792\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.321 [info] [bucket_psi.cc:Run:133] Skip doing psi, because dataset has been aligned!\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.321 [info] [bucket_psi.cc:Run:178] Begin post filtering, indices.size=10792, should_sort=true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.329 [info] [utils.cc:MultiKeySort:88] Executing sort scripts: tail -n +2 data/tmp-sort-in-1684911350322047349 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911350322047349\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.243 [info] [bucket_psi.cc:Init:228] bucket size set to 1048576\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.269 [info] [bucket_psi.cc:Run:97] Begin sanity check for input file: data/bob_train.csv, precheck_switch:true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.281 [info] [csv_checker.cc:CsvChecker:121] Executing duplicated scripts: LC_ALL=C sort --buffer-size=1G --temporary-directory=data --stable selected-keys.1684911350269771195 | LC_ALL=C uniq -d > duplicate-keys.1684911350269771195\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.321 [info] [bucket_psi.cc:Run:115] End sanity check for input file: data/bob_train.csv, size=10792\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.321 [info] [bucket_psi.cc:Run:133] Skip doing psi, because dataset has been aligned!\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.321 [info] [bucket_psi.cc:Run:178] Begin post filtering, indices.size=10792, should_sort=true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.326 [info] [utils.cc:MultiKeySort:88] Executing sort scripts: tail -n +2 data/tmp-sort-in-1684911350321798883 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911350321798883\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.243 [info] [bucket_psi.cc:Init:228] bucket size set to 1048576\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.269 [info] [bucket_psi.cc:Run:97] Begin sanity check for input file: data/alice_train.csv, precheck_switch:true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.288 [info] [csv_checker.cc:CsvChecker:121] Executing duplicated scripts: LC_ALL=C sort --buffer-size=1G --temporary-directory=data --stable selected-keys.1684911350269789933 | LC_ALL=C uniq -d > duplicate-keys.1684911350269789933\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.320 [info] [bucket_psi.cc:Run:115] End sanity check for input file: data/alice_train.csv, size=10792\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.321 [info] [bucket_psi.cc:Run:133] Skip doing psi, because dataset has been aligned!\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.321 [info] [bucket_psi.cc:Run:178] Begin post filtering, indices.size=10792, should_sort=true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.328 [info] [utils.cc:MultiKeySort:88] Executing sort scripts: tail -n +2 data/tmp-sort-in-1684911350321981931 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911350321981931\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.363 [info] [utils.cc:MultiKeySort:90] Finished sort scripts: tail -n +2 data/tmp-sort-in-1684911350322047349 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911350322047349, ret=0\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:50.363 [info] [bucket_psi.cc:Run:216] End post filtering, in=data/carol_train.csv, out=data/carol_train.csv.psi_output_85486\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.357 [info] [utils.cc:MultiKeySort:90] Finished sort scripts: tail -n +2 data/tmp-sort-in-1684911350321798883 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911350321798883, ret=0\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:50.357 [info] [bucket_psi.cc:Run:216] End post filtering, in=data/bob_train.csv, out=data/bob_train.csv.psi_output_85486\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.361 [info] [utils.cc:MultiKeySort:90] Finished sort scripts: tail -n +2 data/tmp-sort-in-1684911350321981931 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911350321981931, ret=0\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:50.362 [info] [bucket_psi.cc:Run:216] End post filtering, in=data/alice_train.csv, out=data/alice_train.csv.psi_output_85486\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.918 [info] [bucket_psi.cc:Init:228] bucket size set to 1048576\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.918 [info] [bucket_psi.cc:Run:97] Begin sanity check for input file: data/carol_test.csv, precheck_switch:true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.925 [info] [csv_checker.cc:CsvChecker:121] Executing duplicated scripts: LC_ALL=C sort --buffer-size=1G --temporary-directory=data --stable selected-keys.1684911351918502537 | LC_ALL=C uniq -d > duplicate-keys.1684911351918502537\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.917 [info] [bucket_psi.cc:Init:228] bucket size set to 1048576\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.917 [info] [bucket_psi.cc:Run:97] Begin sanity check for input file: data/bob_test.csv, precheck_switch:true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.923 [info] [csv_checker.cc:CsvChecker:121] Executing duplicated scripts: LC_ALL=C sort --buffer-size=1G --temporary-directory=data --stable selected-keys.1684911351917485080 | LC_ALL=C uniq -d > duplicate-keys.1684911351917485080\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.917 [info] [bucket_psi.cc:Init:228] bucket size set to 1048576\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.917 [info] [bucket_psi.cc:Run:97] Begin sanity check for input file: data/alice_test.csv, precheck_switch:true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.923 [info] [csv_checker.cc:CsvChecker:121] Executing duplicated scripts: LC_ALL=C sort --buffer-size=1G --temporary-directory=data --stable selected-keys.1684911351917758713 | LC_ALL=C uniq -d > duplicate-keys.1684911351917758713\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.955 [info] [bucket_psi.cc:Run:115] End sanity check for input file: data/carol_test.csv, size=4626\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.956 [info] [bucket_psi.cc:Run:133] Skip doing psi, because dataset has been aligned!\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.956 [info] [bucket_psi.cc:Run:178] Begin post filtering, indices.size=4626, should_sort=true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.959 [info] [utils.cc:MultiKeySort:88] Executing sort scripts: tail -n +2 data/tmp-sort-in-1684911351956618837 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911351956618837\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.986 [info] [utils.cc:MultiKeySort:90] Finished sort scripts: tail -n +2 data/tmp-sort-in-1684911351956618837 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911351956618837, ret=0\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:55:51.986 [info] [bucket_psi.cc:Run:216] End post filtering, in=data/carol_test.csv, out=data/carol_test.csv.psi_output_37067\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.955 [info] [bucket_psi.cc:Run:115] End sanity check for input file: data/bob_test.csv, size=4626\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.956 [info] [bucket_psi.cc:Run:133] Skip doing psi, because dataset has been aligned!\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.956 [info] [bucket_psi.cc:Run:178] Begin post filtering, indices.size=4626, should_sort=true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.957 [info] [utils.cc:MultiKeySort:88] Executing sort scripts: tail -n +2 data/tmp-sort-in-1684911351956284333 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911351956284333\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.983 [info] [utils.cc:MultiKeySort:90] Finished sort scripts: tail -n +2 data/tmp-sort-in-1684911351956284333 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911351956284333, ret=0\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:55:51.983 [info] [bucket_psi.cc:Run:216] End post filtering, in=data/bob_test.csv, out=data/bob_test.csv.psi_output_37067\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.955 [info] [bucket_psi.cc:Run:115] End sanity check for input file: data/alice_test.csv, size=4626\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.956 [info] [bucket_psi.cc:Run:133] Skip doing psi, because dataset has been aligned!\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.956 [info] [bucket_psi.cc:Run:178] Begin post filtering, indices.size=4626, should_sort=true\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.958 [info] [utils.cc:MultiKeySort:88] Executing sort scripts: tail -n +2 data/tmp-sort-in-1684911351956327849 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911351956327849\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.983 [info] [utils.cc:MultiKeySort:90] Finished sort scripts: tail -n +2 data/tmp-sort-in-1684911351956327849 | LC_ALL=C sort --buffer-size=3G --parallel=8 --temporary-directory=./ --stable --field-separator=, --key=1,1 >>data/tmp-sort-out-1684911351956327849, ret=0\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:55:51.983 [info] [bucket_psi.cc:Run:216] End post filtering, in=data/alice_test.csv, out=data/alice_test.csv.psi_output_37067\n",
      "VDataFrame(partitions={alice: Partition(data=<secretflow.device.device.pyu.PYUObject object at 0x7fd000234460>), bob: Partition(data=<secretflow.device.device.pyu.PYUObject object at 0x7fcfb9f9e310>), carol: Partition(data=<secretflow.device.device.pyu.PYUObject object at 0x7fcfb9f9ebb0>)}, aligned=True)\n",
      "Index(['Month', 'WeekOfMonth', 'DayOfWeek', 'Make', 'AccidentArea',\n",
      "       'DayOfWeekClaimed', 'MonthClaimed', 'WeekOfMonthClaimed', 'Sex',\n",
      "       'MaritalStatus', 'Age', 'Fault', 'PolicyType', 'VehicleCategory',\n",
      "       'VehiclePrice', 'PolicyNumber', 'RepNumber', 'Deductible',\n",
      "       'DriverRating', 'Days_Policy_Accident', 'Days_Policy_Claim',\n",
      "       'PastNumberOfClaims', 'AgeOfVehicle', 'AgeOfPolicyHolder',\n",
      "       'PoliceReportFiled', 'WitnessPresent', 'AgentType',\n",
      "       'NumberOfSuppliments', 'AddressChange_Claim', 'NumberOfCars', 'Year',\n",
      "       'BasePolicy', 'FraudFound_P'],\n",
      "      dtype='object')\n"
     ]
    }
   ],
   "source": [
    "\n",
    "from secretflow.data.vertical import read_csv as v_read_csv\n",
    "\n",
    "train_ds = v_read_csv({alice: train_alice_path, bob: train_bob_path, carol: train_carol_path}, keys='UID', drop_keys='UID', spu=my_spu)\n",
    "test_ds = v_read_csv({alice: test_alice_path, bob: test_bob_path, carol: test_carol_path}, keys='UID', drop_keys='UID', spu=my_spu)\n",
    "print(train_ds)\n",
    "print(train_ds.columns)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据洞察\n",
    "基于上层封装的VDataFrame抽象，隐语提供了多种数据分析的API，例如统计信息、查改某些列的信息等。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WeekOfMonth    10792\n",
      "dtype: int64\n",
      "WeekOfMonth    5\n",
      "dtype: int64\n",
      "WeekOfMonth    1\n",
      "dtype: int64\n"
     ]
    }
   ],
   "source": [
    "print(train_ds['WeekOfMonth'].count())\n",
    "\n",
    "print(train_ds['WeekOfMonth'].max())\n",
    "print(train_ds['WeekOfMonth'].min())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据预处理\n",
    "在读取完数据之后，下面我们演示如何在隐语上对一个实际多方持有的数据进行数据预处理。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Label Encoder\n",
    "\n",
    "对无序且二值的值，我们可以使用label encoding，转化为0/1表示"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Col name AccidentArea: ['Urban' 'Rural']\n",
      "Col name Sex: ['Female' 'Male']\n",
      "Col name Fault: ['Policy Holder' 'Third Party']\n",
      "Col name PoliceReportFiled: ['No' 'Yes']\n",
      "Col name WitnessPresent: ['No' 'Yes']\n",
      "Col name AgentType: ['External' 'Internal']\n"
     ]
    }
   ],
   "source": [
    "from secretflow.preprocessing import LabelEncoder\n",
    "\n",
    "cols = ['AccidentArea', 'Sex', 'Fault', 'PoliceReportFiled', 'WitnessPresent', 'AgentType']\n",
    "for col in cols:\n",
    "    print(f\"Col name {col}: {df[col].unique()}\")\n",
    "\n",
    "train_ds_v1 = train_ds.copy()\n",
    "test_ds_v1 = test_ds.copy()\n",
    "\n",
    "label_encoder = LabelEncoder()\n",
    "for col in cols:\n",
    "    label_encoder.fit(train_ds_v1[col])\n",
    "    train_ds_v1[col] = label_encoder.transform(train_ds_v1[col])\n",
    "    test_ds_v1[col] = label_encoder.transform(test_ds_v1[col])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Ordinal) Categorical Features\n",
    "对于有序的类别数据，我们构建映射，将类别数据转化为0~n-1的整数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "orig ds in alice:\n",
      "       Month  WeekOfMonth DayOfWeek       Make  AccidentArea DayOfWeekClaimed  \\\n",
      "0       Jan            1    Friday     Toyota             1        Wednesday   \n",
      "1       Jan            1    Monday    Pontiac             1           Monday   \n",
      "2       Dec            1    Friday     Toyota             1        Wednesday   \n",
      "3       Oct            2    Monday  Chevrolet             1           Monday   \n",
      "4       Sep            4   Tuesday    Pontiac             1          Tuesday   \n",
      "...     ...          ...       ...        ...           ...              ...   \n",
      "10787   Dec            3   Tuesday      Mazda             1        Wednesday   \n",
      "10788   Sep            3   Tuesday       Ford             1        Wednesday   \n",
      "10789   Aug            5  Thursday       Ford             1           Friday   \n",
      "10790   Feb            2  Thursday    Pontiac             1          Tuesday   \n",
      "10791   May            5    Monday  Chevrolet             1          Tuesday   \n",
      "\n",
      "      MonthClaimed  WeekOfMonthClaimed  Sex MaritalStatus  \n",
      "0              Jan                   4    1       Married  \n",
      "1              Jan                   3    1       Married  \n",
      "2              Dec                   2    1       Married  \n",
      "3              Oct                   2    1       Married  \n",
      "4              Sep                   4    0        Single  \n",
      "...            ...                 ...  ...           ...  \n",
      "10787          Dec                   3    1       Married  \n",
      "10788          Oct                   1    1       Married  \n",
      "10789          Sep                   5    1       Married  \n",
      "10790          Feb                   3    1        Single  \n",
      "10791          May                   5    1       Married  \n",
      "\n",
      "[10792 rows x 10 columns]\n",
      "orig ds in alice:\n",
      "        Month  WeekOfMonth  DayOfWeek       Make  AccidentArea  \\\n",
      "0          1            1          5     Toyota             1   \n",
      "1          1            1          1    Pontiac             1   \n",
      "2         12            1          5     Toyota             1   \n",
      "3         10            2          1  Chevrolet             1   \n",
      "4          9            4          2    Pontiac             1   \n",
      "...      ...          ...        ...        ...           ...   \n",
      "10787     12            3          2      Mazda             1   \n",
      "10788      9            3          2       Ford             1   \n",
      "10789      8            5          4       Ford             1   \n",
      "10790      2            2          4    Pontiac             1   \n",
      "10791      5            5          1  Chevrolet             1   \n",
      "\n",
      "       DayOfWeekClaimed  MonthClaimed  WeekOfMonthClaimed  Sex MaritalStatus  \n",
      "0                     3             1                   4    1       Married  \n",
      "1                     1             1                   3    1       Married  \n",
      "2                     3            12                   2    1       Married  \n",
      "3                     1            10                   2    1       Married  \n",
      "4                     2             9                   4    0        Single  \n",
      "...                 ...           ...                 ...  ...           ...  \n",
      "10787                 3            12                   3    1       Married  \n",
      "10788                 3            10                   1    1       Married  \n",
      "10789                 5             9                   5    1       Married  \n",
      "10790                 2             2                   3    1        Single  \n",
      "10791                 2             5                   5    1       Married  \n",
      "\n",
      "[10792 rows x 10 columns]\n"
     ]
    }
   ],
   "source": [
    "cols1 = [\n",
    "    \"Days_Policy_Accident\",\n",
    "    \"Days_Policy_Claim\",\n",
    "    \"AgeOfPolicyHolder\",\n",
    "    \"AddressChange_Claim\",\n",
    "    \"NumberOfCars\",\n",
    "]\n",
    "col_disc = [\n",
    "    {\n",
    "        \"Days_Policy_Accident\": {\n",
    "            \"more than 30\": 31,\n",
    "            \"15 to 30\": 22.5,\n",
    "            \"none\": 0,\n",
    "            \"1 to 7\": 4,\n",
    "            \"8 to 15\": 11.5,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"Days_Policy_Claim\": {\n",
    "            \"more than 30\": 31,\n",
    "            \"15 to 30\": 22.5,\n",
    "            \"8 to 15\": 11.5,\n",
    "            \"none\": 0,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"AgeOfPolicyHolder\": {\n",
    "            \"26 to 30\": 28,\n",
    "            \"31 to 35\": 33,\n",
    "            \"41 to 50\": 45.5,\n",
    "            \"51 to 65\": 58,\n",
    "            \"21 to 25\": 23,\n",
    "            \"36 to 40\": 38,\n",
    "            \"16 to 17\": 16.5,\n",
    "            \"over 65\": 66,\n",
    "            \"18 to 20\": 19,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"AddressChange_Claim\": {\n",
    "            \"1 year\": 1,\n",
    "            \"no change\": 0,\n",
    "            \"4 to 8 years\": 6,\n",
    "            \"2 to 3 years\": 2.5,\n",
    "            \"under 6 months\": 0.5,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"NumberOfCars\": {\n",
    "            \"3 to 4\": 3.5,\n",
    "            \"1 vehicle\": 1,\n",
    "            \"2 vehicles\": 2,\n",
    "            \"5 to 8\": 6.5,\n",
    "            \"more than 8\": 9,\n",
    "        }\n",
    "    },\n",
    "]\n",
    "\n",
    "cols2 = [\n",
    "    \"Month\",\n",
    "    \"DayOfWeek\",\n",
    "    \"DayOfWeekClaimed\",\n",
    "    \"MonthClaimed\",\n",
    "    \"PastNumberOfClaims\",\n",
    "    \"NumberOfSuppliments\",\n",
    "    \"VehiclePrice\",\n",
    "    \"AgeOfVehicle\",\n",
    "]\n",
    "col_ordering = [\n",
    "    {\n",
    "        \"Month\": {\n",
    "            \"Jan\": 1,\n",
    "            \"Feb\": 2,\n",
    "            \"Mar\": 3,\n",
    "            \"Apr\": 4,\n",
    "            \"May\": 5,\n",
    "            \"Jun\": 6,\n",
    "            \"Jul\": 7,\n",
    "            \"Aug\": 8,\n",
    "            \"Sep\": 9,\n",
    "            \"Oct\": 10,\n",
    "            \"Nov\": 11,\n",
    "            \"Dec\": 12,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"DayOfWeek\": {\n",
    "            \"Monday\": 1,\n",
    "            \"Tuesday\": 2,\n",
    "            \"Wednesday\": 3,\n",
    "            \"Thursday\": 4,\n",
    "            \"Friday\": 5,\n",
    "            \"Saturday\": 6,\n",
    "            \"Sunday\": 7,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"DayOfWeekClaimed\": {\n",
    "            \"Monday\": 1,\n",
    "            \"Tuesday\": 2,\n",
    "            \"Wednesday\": 3,\n",
    "            \"Thursday\": 4,\n",
    "            \"Friday\": 5,\n",
    "            \"Saturday\": 6,\n",
    "            \"Sunday\": 7,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"MonthClaimed\": {\n",
    "            \"Jan\": 1,\n",
    "            \"Feb\": 2,\n",
    "            \"Mar\": 3,\n",
    "            \"Apr\": 4,\n",
    "            \"May\": 5,\n",
    "            \"Jun\": 6,\n",
    "            \"Jul\": 7,\n",
    "            \"Aug\": 8,\n",
    "            \"Sep\": 9,\n",
    "            \"Oct\": 10,\n",
    "            \"Nov\": 11,\n",
    "            \"Dec\": 12,\n",
    "        }\n",
    "    },\n",
    "    {\"PastNumberOfClaims\": {\"none\": 0, \"1\": 1, \"2 to 4\": 2, \"more than 4\": 5}},\n",
    "    {\"NumberOfSuppliments\": {\"none\": 0, \"1 to 2\": 1, \"3 to 5\": 3, \"more than 5\": 6}},\n",
    "    {\n",
    "        \"VehiclePrice\": {\n",
    "            \"more than 69000\": 69001,\n",
    "            \"20000 to 29000\": 24500,\n",
    "            \"30000 to 39000\": 34500,\n",
    "            \"less than 20000\": 19999,\n",
    "            \"40000 to 59000\": 49500,\n",
    "            \"60000 to 69000\": 64500,\n",
    "        }\n",
    "    },\n",
    "    {\n",
    "        \"AgeOfVehicle\": {\n",
    "            \"3 years\": 3,\n",
    "            \"6 years\": 6,\n",
    "            \"7 years\": 7,\n",
    "            \"more than 7\": 8,\n",
    "            \"5 years\": 5,\n",
    "            \"new\": 0,\n",
    "            \"4 years\": 4,\n",
    "            \"2 years\": 2,\n",
    "        }\n",
    "    },\n",
    "]\n",
    "\n",
    "from secretflow.data.vertical import VDataFrame\n",
    "\n",
    "def replace(df, col_maps):\n",
    "    df = df.copy()\n",
    "\n",
    "    def func_(df, col_map):\n",
    "        col_name = list(col_map.keys())[0]\n",
    "        col_dict = list(col_map.values())[0]\n",
    "        if col_name not in df.columns:\n",
    "            return\n",
    "        new_list = []\n",
    "        for i in df[col_name]:\n",
    "            new_list.append(col_dict[i])\n",
    "        df[col_name] = new_list\n",
    "\n",
    "    for col_map in col_maps:\n",
    "        func_(df, col_map)\n",
    "    return df\n",
    "\n",
    "col_maps = col_disc + col_ordering\n",
    "\n",
    "train_ds_v2 = train_ds_v1.copy()\n",
    "test_ds_v2 = test_ds_v1.copy()\n",
    "\n",
    "# NOTE: Reveal is only used for demo only!!\n",
    "print(f\"orig ds in alice:\\n {sf.reveal(train_ds_v2.partitions[alice].data)}\")\n",
    "train_ds_v2.partitions[alice].data = alice(replace)(\n",
    "    train_ds_v2.partitions[alice].data, col_maps\n",
    ")\n",
    "train_ds_v2.partitions[bob].data = bob(replace)(\n",
    "    train_ds_v2.partitions[bob].data, col_maps\n",
    ")\n",
    "train_ds_v2.partitions[carol].data = carol(replace)(\n",
    "    train_ds_v2.partitions[carol].data, col_maps\n",
    ")\n",
    "\n",
    "print(f\"orig ds in alice:\\n {sf.reveal(train_ds_v2.partitions[alice].data)}\")\n",
    "\n",
    "test_ds_v2.partitions[alice].data = alice(replace)(\n",
    "    test_ds_v2.partitions[alice].data, col_maps\n",
    ")\n",
    "test_ds_v2.partitions[bob].data = bob(replace)(\n",
    "    test_ds_v2.partitions[bob].data, col_maps\n",
    ")\n",
    "test_ds_v2.partitions[carol].data = carol(replace)(\n",
    "    test_ds_v2.partitions[carol].data, col_maps\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Nominal) Categorical Features\n",
    "无序的类别数据，我们直接采用onehot encoder进行01编码"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Onehot Encoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "orig ds in alice:\n",
      "        Month  WeekOfMonth  DayOfWeek  AccidentArea  DayOfWeekClaimed  \\\n",
      "0          1            1          5             1                 3   \n",
      "1          1            1          1             1                 1   \n",
      "2         12            1          5             1                 3   \n",
      "3         10            2          1             1                 1   \n",
      "4          9            4          2             1                 2   \n",
      "...      ...          ...        ...           ...               ...   \n",
      "10787     12            3          2             1                 3   \n",
      "10788      9            3          2             1                 3   \n",
      "10789      8            5          4             1                 5   \n",
      "10790      2            2          4             1                 2   \n",
      "10791      5            5          1             1                 2   \n",
      "\n",
      "       MonthClaimed  WeekOfMonthClaimed  Sex  Make_Accura  Make_BMW  ...  \\\n",
      "0                 1                   4    1          0.0       0.0  ...   \n",
      "1                 1                   3    1          0.0       0.0  ...   \n",
      "2                12                   2    1          0.0       0.0  ...   \n",
      "3                10                   2    1          0.0       0.0  ...   \n",
      "4                 9                   4    0          0.0       0.0  ...   \n",
      "...             ...                 ...  ...          ...       ...  ...   \n",
      "10787            12                   3    1          0.0       0.0  ...   \n",
      "10788            10                   1    1          0.0       0.0  ...   \n",
      "10789             9                   5    1          0.0       0.0  ...   \n",
      "10790             2                   3    1          0.0       0.0  ...   \n",
      "10791             5                   5    1          0.0       0.0  ...   \n",
      "\n",
      "       Make_Pontiac  Make_Porche  Make_Saab  Make_Saturn  Make_Toyota  \\\n",
      "0               0.0          0.0        0.0          0.0          1.0   \n",
      "1               1.0          0.0        0.0          0.0          0.0   \n",
      "2               0.0          0.0        0.0          0.0          1.0   \n",
      "3               0.0          0.0        0.0          0.0          0.0   \n",
      "4               1.0          0.0        0.0          0.0          0.0   \n",
      "...             ...          ...        ...          ...          ...   \n",
      "10787           0.0          0.0        0.0          0.0          0.0   \n",
      "10788           0.0          0.0        0.0          0.0          0.0   \n",
      "10789           0.0          0.0        0.0          0.0          0.0   \n",
      "10790           1.0          0.0        0.0          0.0          0.0   \n",
      "10791           0.0          0.0        0.0          0.0          0.0   \n",
      "\n",
      "       Make_VW  MaritalStatus_Divorced  MaritalStatus_Married  \\\n",
      "0          0.0                     0.0                    1.0   \n",
      "1          0.0                     0.0                    1.0   \n",
      "2          0.0                     0.0                    1.0   \n",
      "3          0.0                     0.0                    1.0   \n",
      "4          0.0                     0.0                    0.0   \n",
      "...        ...                     ...                    ...   \n",
      "10787      0.0                     0.0                    1.0   \n",
      "10788      0.0                     0.0                    1.0   \n",
      "10789      0.0                     0.0                    1.0   \n",
      "10790      0.0                     0.0                    0.0   \n",
      "10791      0.0                     0.0                    1.0   \n",
      "\n",
      "       MaritalStatus_Single  MaritalStatus_Widow  \n",
      "0                       0.0                  0.0  \n",
      "1                       0.0                  0.0  \n",
      "2                       0.0                  0.0  \n",
      "3                       0.0                  0.0  \n",
      "4                       1.0                  0.0  \n",
      "...                     ...                  ...  \n",
      "10787                   0.0                  0.0  \n",
      "10788                   0.0                  0.0  \n",
      "10789                   0.0                  0.0  \n",
      "10790                   1.0                  0.0  \n",
      "10791                   0.0                  0.0  \n",
      "\n",
      "[10792 rows x 31 columns]\n"
     ]
    }
   ],
   "source": [
    "from secretflow.preprocessing import OneHotEncoder\n",
    "onehot_cols = ['Make','MaritalStatus','PolicyType','VehicleCategory','BasePolicy']\n",
    "\n",
    "onehot_encoder = OneHotEncoder()\n",
    "onehot_encoder.fit(train_ds_v2[onehot_cols])\n",
    "\n",
    "enc_feats = onehot_encoder.transform(train_ds_v2[onehot_cols])\n",
    "feature_names = enc_feats.columns\n",
    "train_ds_v3 = train_ds_v2.drop(columns=onehot_cols)\n",
    "train_ds_v3[feature_names] = enc_feats\n",
    "\n",
    "\n",
    "enc_feats = onehot_encoder.transform(test_ds_v2[onehot_cols])\n",
    "test_ds_v3 = test_ds_v2.drop(columns=onehot_cols)\n",
    "test_ds_v3[feature_names] = enc_feats\n",
    "\n",
    "print(f\"orig ds in alice:\\n {sf.reveal(train_ds_v3.partitions[alice].data)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data load done\n"
     ]
    }
   ],
   "source": [
    "train_ds_final = train_ds_v3.copy()\n",
    "test_ds_final = test_ds_v3.copy()\n",
    "\n",
    "X_train = train_ds_v3.drop(columns=['FraudFound_P'])\n",
    "y_train = train_ds_final['FraudFound_P']\n",
    "X_test = test_ds_final.drop(columns='FraudFound_P')\n",
    "y_test = test_ds_final['FraudFound_P']\n",
    "\n",
    "print(\"data load done\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据对象转换\n",
    "此处我们将PYUObject 转化为 SPUObject，方便输入到SPU device执行基于MPC协议的隐私计算"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=2294419)\u001b[0m [2023-05-24 14:55:57.641] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(_run pid=2294389)\u001b[0m [2023-05-24 14:55:57.703] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(_run pid=2294392)\u001b[0m [2023-05-24 14:55:57.670] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=2294392)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: \n",
      "\u001b[2m\u001b[36m(_run pid=2294392)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294392)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294392)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=2294392)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=2294392)\u001b[0m WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n",
      "\u001b[2m\u001b[36m(_run pid=2294389)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: \n",
      "\u001b[2m\u001b[36m(_run pid=2294389)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294389)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294389)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=2294389)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=2294389)\u001b[0m WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=2294562)\u001b[0m [2023-05-24 14:55:57.944] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "X_train type: VDataFrame(partitions={alice: Partition(data=<secretflow.device.device.pyu.PYUObject object at 0x7fcfb9f7b0a0>), bob: Partition(data=<secretflow.device.device.pyu.PYUObject object at 0x7fcfb9f7b070>), carol: Partition(data=<secretflow.device.device.pyu.PYUObject object at 0x7fcfb9f61910>)}, aligned=True)\n",
      "\n",
      "X_train_spu type: <secretflow.device.device.spu.SPUObject object at 0x7fcfb9f07610>\n",
      "[2023-05-24 14:55:58.124] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "X_train_plaintext: \n",
      "[[ 1.  1.  5. ...  1.  0.  0.]\n",
      " [ 1.  1.  1. ...  0.  0.  1.]\n",
      " [12.  1.  5. ...  0.  1.  0.]\n",
      " ...\n",
      " [ 8.  5.  4. ...  1.  0.  0.]\n",
      " [ 2.  2.  4. ...  0.  1.  0.]\n",
      " [ 5.  5.  1. ...  0.  1.  0.]]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_run pid=2294562)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: \n",
      "\u001b[2m\u001b[36m(_run pid=2294562)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294562)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "\u001b[2m\u001b[36m(_run pid=2294562)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "\u001b[2m\u001b[36m(_run pid=2294562)\u001b[0m INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "\u001b[2m\u001b[36m(_run pid=2294562)\u001b[0m WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "\n",
    "\"\"\"\n",
    "Convert the VDataFrame object to SPUObject\n",
    "\"\"\"\n",
    "def vdataframe_to_spu(vdf: VDataFrame):\n",
    "    spu_partitions = []\n",
    "    for device in [alice, bob, carol]:\n",
    "        spu_partitions.append(vdf.partitions[device].data.to(my_spu))\n",
    "    base_partition = spu_partitions[0]\n",
    "    for i in range(1, len(spu_partitions)):\n",
    "        base_partition = my_spu(lambda x, y: jnp.concatenate([x, y], axis=1))(\n",
    "            base_partition, spu_partitions[i]\n",
    "        )\n",
    "    return base_partition\n",
    "\n",
    "\n",
    "X_train_spu = vdataframe_to_spu(X_train)\n",
    "y_train_spu = y_train.partitions[carol].data.to(my_spu)\n",
    "X_test_spu = vdataframe_to_spu(X_test)\n",
    "y_test_spu = y_test.partitions[carol].data.to(my_spu)\n",
    "print(f\"X_train type: {X_train}\\n\\nX_train_spu type: {X_train_spu}\")\n",
    "\n",
    "\"\"\"\n",
    "NOTE: This is only for demo only!! This shall not be used in production.\n",
    "\"\"\"\n",
    "X_train_plaintext = sf.reveal(X_train_spu)\n",
    "y_train_plaintext = sf.reveal(y_train_spu)\n",
    "X_test_plaintext = sf.reveal(X_test_spu)\n",
    "y_test_plaintext = sf.reveal(y_test_spu)\n",
    "\n",
    "print(f'X_train_plaintext: \\n{X_train_plaintext}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型构建\n",
    "在完成数据的读入之后，下面我们进行模型的构建。在本demo中，主要提供了三种模型的构建：\n",
    "- LR: 逻辑回归\n",
    "- NN：神经网络模型\n",
    "- XGB: XGBoost 树模型\n",
    "\n",
    "> 注意，本示例主要是演示在隐语上进行算法开发的流程，并没有针对模型 (LR, NN) 进行调参。我们分别提供了明文和密文的计算结果，实验结果显示两者的输出是基本一致的，表明隐语的密态计算能够和明文计算保持精度一致。"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LR ( jax ) using SPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.example_libraries import optimizers, stax\n",
    "from jax.example_libraries.stax import (\n",
    "    Conv,\n",
    "    MaxPool,\n",
    "    AvgPool,\n",
    "    Flatten,\n",
    "    Dense,\n",
    "    Relu,\n",
    "    Sigmoid,\n",
    "    LogSoftmax,\n",
    "    Softmax,\n",
    "    BatchNorm,\n",
    ")\n",
    "\n",
    "\n",
    "def sigmoid(x):\n",
    "    x = (x - jnp.min(x)) / (jnp.max(x) - jnp.min(x))\n",
    "    return 1 / (1 + jnp.exp(-x))\n",
    "\n",
    "# Outputs probability of a label being true.\n",
    "def predict_lr(W, b, inputs):\n",
    "    return sigmoid(jnp.dot(inputs, W) + b)\n",
    "\n",
    "# Training loss is the negative log-likelihood of the training examples.\n",
    "def loss_lr(W, b, inputs, targets):\n",
    "    preds = predict_lr(W, b, inputs)\n",
    "    label_probs = preds * targets + (1 - preds) * (1 - targets)\n",
    "    return -jnp.mean(jnp.log(label_probs))\n",
    "\n",
    "def train_step(W, b, X, y, learning_rate):\n",
    "    loss_value, Wb_grad = jax.value_and_grad(loss_lr, (0, 1))(W, b, X, y)\n",
    "    W -= learning_rate * Wb_grad[0]\n",
    "    b -= learning_rate * Wb_grad[1]\n",
    "    return loss_value, W, b\n",
    "\n",
    "def fit(W, b, X, y, epochs=1, learning_rate=1e-2, batch_size=128):\n",
    "    losses = jnp.array([])\n",
    "\n",
    "    xs = jnp.array_split(X, len(X) / batch_size, axis=0)\n",
    "    ys = jnp.array_split(y, len(y) / batch_size, axis=0)\n",
    "\n",
    "    for _ in range(epochs):\n",
    "        for (batch_x, batch_y) in zip(xs, ys):\n",
    "            l, W, b = train_step(\n",
    "                W, b, batch_x, batch_y, learning_rate=learning_rate\n",
    "            )\n",
    "            losses = jnp.append(losses, l)\n",
    "    return losses, W, b\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu_driver': NOT_FOUND: Unable to find driver in registry given worker: \n",
      "INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'cuda': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'\n",
      "INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'tpu': INVALID_ARGUMENT: TpuPlatform is not available.\n",
      "INFO:jax._src.lib.xla_bridge:Unable to initialize backend 'plugin': xla_extension has no attributes named get_plugin_device_client. Compile TensorFlow with //tensorflow/compiler/xla/python:enable_plugin_device set to true (defaults to false) to enable this.\n",
      "WARNING:jax._src.lib.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31m(Jax LR CPU) auc: 0.524559290927313\u001b[0m\n",
      "\u001b[31m(Jax LR SPU) auc: 0.48416274742946\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from jax import random\n",
    "import sys\n",
    "import time\n",
    "import logging\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)\n",
    "logging.getLogger().setLevel(logging.INFO)\n",
    "\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "\n",
    "# Hyperparameter\n",
    "key = random.PRNGKey(42)\n",
    "W = jax.random.normal(key, shape=(64,))\n",
    "b = 0.0\n",
    "epochs = 1\n",
    "learning_rate = 1e-2\n",
    "batch_size = 128\n",
    "\n",
    "\"\"\"\n",
    "CPU-version plaintext computation\n",
    "\"\"\"\n",
    "losses_cpu, W_cpu, b_cpu = fit(\n",
    "    W,\n",
    "    b,\n",
    "    X_train_plaintext,\n",
    "    y_train_plaintext,\n",
    "    epochs=epochs,\n",
    "    learning_rate=learning_rate,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "y_pred_cpu = predict_lr(W_cpu, b_cpu, X_test_plaintext)\n",
    "print(f\"\\033[31m(Jax LR CPU) auc: {roc_auc_score(y_test_plaintext, y_pred_cpu)}\\033[0m\")\n",
    "\n",
    "\"\"\"\n",
    "SPU-version secure computation\n",
    "\"\"\"\n",
    "W_, b_ = (\n",
    "    sf.to(alice, W).to(my_spu),\n",
    "    sf.to(alice, b).to(my_spu),\n",
    ")\n",
    "losses_spu, W_spu, b_spu = my_spu(\n",
    "    fit,\n",
    "    static_argnames=[\"epochs\", \"learning_rate\", \"batch_size\"],\n",
    "    num_returns_policy=sf.device.SPUCompilerNumReturnsPolicy.FROM_USER,\n",
    "    user_specified_num_returns=3,\n",
    ")(\n",
    "    W_,\n",
    "    b_,\n",
    "    X_train_spu,\n",
    "    y_train_spu,\n",
    "    epochs=epochs,\n",
    "    learning_rate=learning_rate,\n",
    "    batch_size=batch_size,\n",
    ")\n",
    "\n",
    "y_pred_spu = my_spu(predict_lr)(W_spu, b_spu, X_test_spu)\n",
    "y_pred = sf.reveal(y_pred_spu)\n",
    "print(f\"\\033[31m(Jax LR SPU) auc: {roc_auc_score(y_test_plaintext, y_pred)}\\033[0m\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NN ( jax + flax ) using SPU"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mWARNING: There was an error checking the latest version of pip.\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "!{sys.executable} -m pip install flax==0.6.0 -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31m(Flax NN CPU) auc: 0.5022025986877813\u001b[0m\n",
      "\u001b[31m(Flax NN SPU) auc: 0.5022034401772514\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from typing import Sequence\n",
    "import flax.linen as nn\n",
    "\n",
    "\n",
    "class MLP(nn.Module):\n",
    "    features: Sequence[int]\n",
    "\n",
    "    @nn.compact\n",
    "    def __call__(self, x):\n",
    "        for feat in self.features[:-1]:\n",
    "            x = nn.relu(nn.Dense(feat)(x))\n",
    "        x = nn.Dense(self.features[-1])(x)\n",
    "        return x\n",
    "\n",
    "FEATURES = [1]\n",
    "flax_nn = MLP(FEATURES)\n",
    "\n",
    "def predict(params, x):\n",
    "    from typing import Sequence\n",
    "    import flax.linen as nn\n",
    "    class MLP(nn.Module):\n",
    "        features: Sequence[int]\n",
    "\n",
    "        @nn.compact\n",
    "        def __call__(self, x):\n",
    "            for feat in self.features[:-1]:\n",
    "                x = nn.relu(nn.Dense(feat)(x))\n",
    "            x = nn.Dense(self.features[-1])(x)\n",
    "            return x\n",
    "    FEATURES = [1]\n",
    "    flax_nn = MLP(FEATURES)\n",
    "    return flax_nn.apply(params, x)\n",
    "\n",
    "\n",
    "def loss_func(params, x, y):\n",
    "    preds = predict(params, x)\n",
    "    label_probs = preds * y + (1 - preds) * (1 - y)\n",
    "    return -jnp.mean(jnp.log(label_probs))\n",
    "\n",
    "\n",
    "def train_auto_grad(X, y, params, batch_size=10, epochs=10, learning_rate=0.01):\n",
    "    xs = jnp.array_split(X, len(X) / batch_size, axis=0)\n",
    "    ys = jnp.array_split(y, len(y) / batch_size, axis=0)\n",
    "\n",
    "    for _ in range(epochs):\n",
    "        for (batch_x, batch_y) in zip(xs, ys):\n",
    "            _, grads = jax.value_and_grad(loss_func)(params, batch_x, batch_y)\n",
    "            params = jax.tree_util.tree_map(\n",
    "                lambda p, g: p - learning_rate * g, params, grads\n",
    "            )\n",
    "    return params\n",
    "\n",
    "epochs = 1\n",
    "learning_rate = 1e-2\n",
    "batch_size = 128\n",
    "\n",
    "feature_dim = 64  # from the dataset\n",
    "init_params = flax_nn.init(jax.random.PRNGKey(1), jnp.ones((batch_size, feature_dim)))\n",
    "\n",
    "\"\"\"\n",
    "CPU-version plaintext computation\n",
    "\"\"\"\n",
    "params = train_auto_grad(\n",
    "    X_train_plaintext, y_train_plaintext, init_params, batch_size, epochs, learning_rate\n",
    ")\n",
    "y_pred = predict(params, X_test_plaintext)\n",
    "print(f\"\\033[31m(Flax NN CPU) auc: {roc_auc_score(y_test_plaintext, y_pred)}\\033[0m\")\n",
    "\n",
    "\"\"\"\n",
    "SPU-version secure computation\n",
    "\"\"\"\n",
    "params_spu = sf.to(alice, init_params).to(my_spu)\n",
    "params_spu = my_spu(train_auto_grad, static_argnames=['batch_size', 'epochs', 'learning_rate'])(\n",
    "    X_train_spu, y_train_spu, params_spu, batch_size=batch_size, epochs=epochs, learning_rate=learning_rate\n",
    ")\n",
    "y_pred_spu = my_spu(predict)(params_spu, X_test_spu)\n",
    "y_pred_ = sf.reveal(y_pred_spu)\n",
    "print(f\"\\033[31m(Flax NN SPU) auc: {roc_auc_score(y_test_plaintext, y_pred_)}\\033[0m\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### XGB ( jax ) using SPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:Create proxy actor <class 'secretflow.ml.boost.ss_xgb_v.core.tree_worker.XgbTreeWorker'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.boost.ss_xgb_v.core.tree_worker.XgbTreeWorker'> with party bob.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.boost.ss_xgb_v.core.tree_worker.XgbTreeWorker'> with party carol.\n",
      "INFO:root:fragment_count 1\n",
      "INFO:root:prepare time 0.3021812438964844s\n",
      "INFO:root:global_setup time 2.1790771484375s\n",
      "INFO:root:build & infeed bucket_map fragments [0, 0]\n",
      "INFO:root:build & infeed bucket_map time 0.32753467559814453s\n",
      "INFO:root:init_pred time 0.04900479316711426s\n",
      "INFO:root:epoch 0 tree_setup time 0.11261296272277832s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_spu_compile pid=2294562)\u001b[0m /* error: missing value */\n",
      "\u001b[2m\u001b[36m(_spu_compile pid=2294562)\u001b[0m {}:task_name:_spu_compile\n",
      "\u001b[2m\u001b[36m(_spu_compile pid=2294562)\u001b[0m /* error: missing value */\n",
      "\u001b[2m\u001b[36m(_spu_compile pid=2294562)\u001b[0m {}:task_name:_spu_compile\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:fragment[0, 0] gradient sum time 0.46968507766723633s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(SPURuntime pid=2300421)\u001b[0m 2023-05-24 14:57:37.774 [info] [thread_pool.cc:ThreadPool:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300420)\u001b[0m 2023-05-24 14:57:37.772 [info] [thread_pool.cc:ThreadPool:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(SPURuntime pid=2300419)\u001b[0m 2023-05-24 14:57:37.777 [info] [thread_pool.cc:ThreadPool:30] Create a fixed thread pool with size 63\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:level 0 time 0.8138599395751953s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.42415881156921387s\n",
      "INFO:root:level 1 time 0.7801706790924072s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.5980944633483887s\n",
      "INFO:root:level 2 time 1.0563251972198486s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.5459139347076416s\n",
      "INFO:root:level 3 time 0.9727473258972168s\n",
      "INFO:root:epoch 0 time 3.961634874343872s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(XgbTreeWorker pid=2312009)\u001b[0m [2023-05-24 14:57:41.177] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(XgbTreeWorker pid=2312010)\u001b[0m [2023-05-24 14:57:41.195] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(XgbTreeWorker pid=2312013)\u001b[0m [2023-05-24 14:57:41.212] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:epoch 1 tree_setup time 0.1868457794189453s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_spu_compile pid=2294389)\u001b[0m /* error: missing value */\n",
      "\u001b[2m\u001b[36m(_spu_compile pid=2294389)\u001b[0m {}:task_name:_spu_compile\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:fragment[0, 0] gradient sum time 0.4389340877532959s\n",
      "INFO:root:level 0 time 0.7621815204620361s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4361765384674072s\n",
      "INFO:root:level 1 time 0.7648963928222656s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.46132969856262207s\n",
      "INFO:root:level 2 time 0.8321614265441895s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.5284152030944824s\n",
      "INFO:root:level 3 time 0.9844522476196289s\n",
      "INFO:root:epoch 1 time 3.611001968383789s\n",
      "INFO:root:epoch 2 tree_setup time 0.11684489250183105s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.3915536403656006s\n",
      "INFO:root:level 0 time 0.7235503196716309s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.42185139656066895s\n",
      "INFO:root:level 1 time 0.7636382579803467s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.46842050552368164s\n",
      "INFO:root:level 2 time 0.8526918888092041s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.5360352993011475s\n",
      "INFO:root:level 3 time 0.9925544261932373s\n",
      "INFO:root:epoch 2 time 3.6581156253814697s\n",
      "INFO:root:epoch 3 tree_setup time 0.17425537109375s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_spu_compile pid=2294392)\u001b[0m /* error: missing value */\n",
      "\u001b[2m\u001b[36m(_spu_compile pid=2294392)\u001b[0m {}:task_name:_spu_compile\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:fragment[0, 0] gradient sum time 0.4217956066131592s\n",
      "INFO:root:level 0 time 0.7510805130004883s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4604473114013672s\n",
      "INFO:root:level 1 time 0.9107728004455566s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.46227431297302246s\n",
      "INFO:root:level 2 time 0.8462221622467041s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.5579540729522705s\n",
      "INFO:root:level 3 time 1.005491018295288s\n",
      "INFO:root:epoch 3 time 3.8129096031188965s\n",
      "INFO:root:epoch 4 tree_setup time 0.16765856742858887s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2m\u001b[36m(_spu_compile pid=2294419)\u001b[0m /* error: missing value */\n",
      "\u001b[2m\u001b[36m(_spu_compile pid=2294419)\u001b[0m {}:task_name:_spu_compile\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:root:fragment[0, 0] gradient sum time 0.42427778244018555s\n",
      "INFO:root:level 0 time 0.6625769138336182s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4184889793395996s\n",
      "INFO:root:level 1 time 0.7157876491546631s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4058997631072998s\n",
      "INFO:root:level 2 time 0.7082300186157227s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.5008931159973145s\n",
      "INFO:root:level 3 time 0.8738894462585449s\n",
      "INFO:root:epoch 4 time 3.195892333984375s\n",
      "INFO:root:epoch 5 tree_setup time 0.14134454727172852s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.400066614151001s\n",
      "INFO:root:level 0 time 0.641761064529419s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.3746778964996338s\n",
      "INFO:root:level 1 time 0.6471817493438721s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.42211461067199707s\n",
      "INFO:root:level 2 time 0.72641921043396s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.46430206298828125s\n",
      "INFO:root:level 3 time 0.8361771106719971s\n",
      "INFO:root:epoch 5 time 3.0644524097442627s\n",
      "INFO:root:epoch 6 tree_setup time 0.11562013626098633s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.3859133720397949s\n",
      "INFO:root:level 0 time 0.635554313659668s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.41338205337524414s\n",
      "INFO:root:level 1 time 0.6643311977386475s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.3950026035308838s\n",
      "INFO:root:level 2 time 0.6798880100250244s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.45518994331359863s\n",
      "INFO:root:level 3 time 0.8225488662719727s\n",
      "INFO:root:epoch 6 time 3.0115859508514404s\n",
      "INFO:root:epoch 7 tree_setup time 0.11349105834960938s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.37856125831604004s\n",
      "INFO:root:level 0 time 0.6111841201782227s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4059903621673584s\n",
      "INFO:root:level 1 time 0.7024598121643066s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.421400785446167s\n",
      "INFO:root:level 2 time 0.7304990291595459s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4517638683319092s\n",
      "INFO:root:level 3 time 0.8225312232971191s\n",
      "INFO:root:epoch 7 time 3.069700002670288s\n",
      "INFO:root:epoch 8 tree_setup time 0.1077120304107666s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.3767871856689453s\n",
      "INFO:root:level 0 time 0.6241645812988281s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.374847412109375s\n",
      "INFO:root:level 1 time 0.6185367107391357s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.40407586097717285s\n",
      "INFO:root:level 2 time 0.7317712306976318s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.45026302337646484s\n",
      "INFO:root:level 3 time 0.8427674770355225s\n",
      "INFO:root:epoch 8 time 3.035346269607544s\n",
      "INFO:root:epoch 9 tree_setup time 0.13258838653564453s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4179399013519287s\n",
      "INFO:root:level 0 time 0.6501905918121338s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.38472986221313477s\n",
      "INFO:root:level 1 time 0.6701993942260742s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.4303562641143799s\n",
      "INFO:root:level 2 time 0.7147767543792725s\n",
      "INFO:root:fragment[0, 0] gradient sum time 0.46015405654907227s\n",
      "INFO:root:level 3 time 0.8475117683410645s\n",
      "INFO:root:epoch 9 time 3.0256881713867188s\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.boost.ss_xgb_v.core.tree_worker.XgbTreeWorker'> with party alice.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.boost.ss_xgb_v.core.tree_worker.XgbTreeWorker'> with party bob.\n",
      "INFO:root:Create proxy actor <class 'secretflow.ml.boost.ss_xgb_v.core.tree_worker.XgbTreeWorker'> with party carol.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train time: 37.70984649658203\n",
      "\u001b[2m\u001b[36m(XgbTreeWorker pid=2316828)\u001b[0m [2023-05-24 14:58:14.123] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(XgbTreeWorker pid=2316834)\u001b[0m [2023-05-24 14:58:14.274] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "\u001b[2m\u001b[36m(XgbTreeWorker pid=2316829)\u001b[0m [2023-05-24 14:58:14.254] [info] [thread_pool.cc:30] Create a fixed thread pool with size 63\n",
      "predict time: 3.1688623428344727\n",
      "\u001b[31m(SS-XGB) auc: 0.8239633480846438\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "from secretflow.ml.boost.ss_xgb_v import Xgb\n",
    "import time\n",
    "from sklearn.metrics import roc_auc_score\n",
    "\n",
    "\"\"\"\n",
    "SPU-version Secure computation\n",
    "\"\"\"\n",
    "xgb = Xgb(my_spu)\n",
    "params = {\n",
    "    # <<< !!! >>> change args to your test settings.\n",
    "    # for more detail, see Xgb.train.__doc__\n",
    "    'num_boost_round': 10,\n",
    "    'max_depth': 4,\n",
    "    'learning_rate': 0.05,\n",
    "    'sketch_eps': 0.05,\n",
    "    'objective': 'logistic',\n",
    "    'reg_lambda': 1,\n",
    "    'subsample': 0.75,\n",
    "    'colsample_by_tree': 1,\n",
    "    'base_score': 0.5,\n",
    "}\n",
    "\n",
    "start = time.time()\n",
    "model = xgb.train(params, X_train, y_train)\n",
    "print(f\"train time: {time.time() - start}\")\n",
    "\n",
    "start =time.time()\n",
    "spu_yhat = model.predict(X_test)\n",
    "print(f\"predict time: {time.time() - start}\")\n",
    "\n",
    "yhat = sf.reveal(spu_yhat)\n",
    "print(f\"\\033[31m(SS-XGB) auc: {roc_auc_score(y_test_plaintext, yhat)}\\033[0m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[14:58:15] WARNING: ../src/learner.cc:767: \n",
      "Parameters: { \"n_estimators\", \"sketch_eps\" } are not used.\n",
      "\n",
      "\u001b[31m(Sklearn-XGB) auc: 0.8231883362827539\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "Plaintext baseline\n",
    "\"\"\"\n",
    "import xgboost as SKxgb\n",
    "\n",
    "params = {\n",
    "    # <<< !!! >>> change args to your test settings.\n",
    "    # for more detail, see Xgb.train.__doc__\n",
    "    \"n_estimators\": 10,\n",
    "    \"max_depth\": 4,\n",
    "    'eval_metric': 'auc',\n",
    "    \"learning_rate\": 0.05,\n",
    "    \"sketch_eps\": 0.05,\n",
    "    \"objective\": \"binary:logistic\",\n",
    "    \"reg_lambda\": 1,\n",
    "    \"subsample\": 0.75,\n",
    "    \"colsample_by_tree\": 1,\n",
    "    \"base_score\": 0.5,\n",
    "}\n",
    "dtrain = SKxgb.DMatrix(X_train_plaintext, label=y_train_plaintext)\n",
    "bst = SKxgb.train(params, dtrain, params[\"n_estimators\"])\n",
    "dtest = SKxgb.DMatrix(X_test_plaintext)\n",
    "y_pred = bst.predict(dtest)\n",
    "print(f\"\\033[31m(Sklearn-XGB) auc: {roc_auc_score(y_test_plaintext, y_pred)}\\033[0m\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The End\n",
    "显示地调用`sf.shutdown()`关闭实例化的集群。\n",
    "> 注意：如果是在.py文件中运行代码，不需要显示地执行shutdown，在程序进程运行结束后会隐式地执行`shutdown`函数。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "sf.shutdown()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 小结一下\n",
    "- 介绍了如何针对一个实际场景的应用，在隐语上进行开发，提供隐私保护的能力\n",
    "- 隐语上的数据加载、预处理、建模、训练流程\n",
    "- 下一步，自己实现任意的计算（jax实现的计算），对于TF，pytorch的支持WIP"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "sf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "a8435d1c2867c43d8e4505e2a7d45b70fc4003dc549050cc77e922e23eb0b16c"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
